import cv2
import easyocr
import numpy as np
import Levenshtein

reader = easyocr.Reader(['en'], gpu=False)  # 支持 ['ch_sim', 'en'] 等多语言

def extract_easyocr_text(image):
    results = reader.readtext(image)
    boxes = []
    for (bbox, text, conf) in results:
        if len(text.strip()) == 0 or conf < 0.5:
            continue
        (x1, y1), (_, _), (x2, y2), (_, _) = bbox
        boxes.append({'text': text.strip(), 'bbox': (int(x1), int(y1), int(x2), int(y2))})
    return boxes

def compare_texts(texts1, texts2, threshold=0.8):
    diffs1, diffs2 = [], []
    for t1 in texts1:
        best_score = 0
        for t2 in texts2:
            score = Levenshtein.ratio(t1['text'], t2['text'])
            best_score = max(best_score, score)
        if best_score < threshold:
            diffs1.append(t1['bbox'])
    for t2 in texts2:
        best_score = 0
        for t1 in texts1:
            score = Levenshtein.ratio(t2['text'], t1['text'])
            best_score = max(best_score, score)
        if best_score < threshold:
            diffs2.append(t2['bbox'])
    return diffs1, diffs2

def highlight_differences(image, boxes):
    for (x1, y1, x2, y2) in boxes:
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 2)
    return image

def compare_images(img_path1, img_path2, output_path='diff_output.png'):
    img1 = cv2.imread(img_path1)
    img2 = cv2.imread(img_path2)

    height = max(img1.shape[0], img2.shape[0])
    width = max(img1.shape[1], img2.shape[1])
    img1 = cv2.resize(img1, (width, height))
    img2 = cv2.resize(img2, (width, height))

    texts1 = extract_easyocr_text(img1)
    texts2 = extract_easyocr_text(img2)

    diffs1, diffs2 = compare_texts(texts1, texts2)

    img1_marked = highlight_differences(img1.copy(), diffs1)
    img2_marked = highlight_differences(img2.copy(), diffs2)

    combined = np.hstack((img1_marked, img2_marked))
    cv2.imwrite(output_path, combined)
    print(f"[✓] 输出图保存至：{output_path}")

if __name__ == "__main__":
    compare_images("compare_1.png", "compare_2.png", "diff_output.png")
